import xgboost as xgb
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import root_mean_squared_error

"""
XGBoost 二分类 预测是否为“virginica”
"""


X, y = load_iris(return_X_y=True)

# 展示数据类型
print("X的数据类型:", type(X))  # numpy.ndarray
print("y的数据类型:", type(y))  # numpy.ndarray
print("X的形状:", X.shape)
print("y的形状:", y.shape)

# 展示数据集信息
print("数据集特征数:", X.shape[1])
print("数据集样本数:", X.shape[0])
print("y的前三个值", y[:3])


# 划分数据集
# train_test_split() 函数将数据集划分为训练集和测试集
# test_size=0.3 表示测试集占总数据的30%，训练集占70%
# random_state=42 设置随机种子，确保每次划分结果一致，便于结果复现
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.3, random_state=42
)


def train_model(X_train, y_train, max_depth=3, n_estimators=100):
    """
    训练 XGBoost 模型
    :param X_train: 训练特征数据
    :param y_train: 训练目标数据
    :param max_depth: 树的最大深度
    :param n_estimators: 树的数量
    :return: 训练好的模型实例
    """
    clf = xgb.XGBClassifier(
        objective="binary:logistic", max_depth=max_depth, n_estimators=n_estimators
    )
    clf.fit(X_train, y_train)
    return clf


def evaluate_model(model, X_test, y_test):
    """
    评估模型性能
    :param model: 训练好的模型实例
    :param X_test: 测试集特征
    :param y_test: 测试集目标值
    :return: 均方根误差（RMSE）
    """
    # 使用测试集进行预测，并计算均方根误差（RMSE）
    rmse = root_mean_squared_error(y_test, model.predict(X_test))
    # RMSE是回归问题常用的评估指标，值越小表示模型预测越准确
    print("RMSE:", rmse)
    return rmse


def save_model(model, model_path):
    """
    保存模型到指定路径
    :param model: 训练好的模型实例
    :param model_path: 模型保存路径
    """
    if not model_path.endswith(".model"):
        model_path += ".model"
    model.save_model(model_path)
    print(f"模型已保存至 {model_path}")


def load_model(model_path):
    """
    从指定路径加载模型
    :param model_path: 模型保存路径
    :return: 加载的模型实例
    """
    model = xgb.XGBRegressor()
    model.load_model(model_path)
    print(f"模型已从 {model_path} 加载")
    return model


def predict(model, X):
    """
    使用模型进行预测
    :param model: 训练好的模型实例
    :param X: 输入特征数据
    :return: 预测结果
    """
    return model.predict(X)


# 主程序执行
if __name__ == "__main__":
    # 训练模型
    print("开始训练模型...")
    model = train_model(X_train, y_train)

    # 评估模型
    print("模型评估结果:")
    evaluate_model(model, X_test, y_test)

    # 保存模型
    # save_model(model, "diabetes_xgboost_model.json")

    # 加载模型
    # print("\n验证模型加载功能:")
    # loaded_model = load_model("diabetes_xgboost_model.json")

    # 使用模型，预测值是多少
